"""生成智能体在 Overcooked 环境中执行的 GIF 动画"""
import argparse
import os
import json
import torch
import numpy as np

try:
    import imageio
    HAS_IMAGEIO = True
except ImportError:
    HAS_IMAGEIO = False
    print("Warning: imageio not installed. Please run: pip install imageio")

from src.envs.overcooked.overcooked_env import Overcooked
from src.algo.r_mappo.rMAPPOPolicy import R_MAPPOPolicy


def parse_args():
    parser = argparse.ArgumentParser(description='生成 Overcooked 评估 GIF')

    parser.add_argument('--model_dir', type=str, required=True,
                        help='模型检查点目录')
    parser.add_argument('--layout_name', type=str, default='small_coor_2',
                        help='Overcooked 地图名称')
    parser.add_argument('--n_episodes', type=int, default=3,
                        help='生成多少个 episode 的 GIF')
    parser.add_argument('--output_dir', type=str, default='renders',
                        help='输出目录')
    parser.add_argument('--fps', type=int, default=10,
                        help='GIF 帧率（每秒帧数）')
    parser.add_argument('--max_steps', type=int, default=300,
                        help='每个 episode 的最大步数')
    parser.add_argument('--deterministic', action='store_true',
                        help='使用确定性策略')
    parser.add_argument('--cuda', action='store_true', default=False)

    return parser.parse_args()


def load_policies(model_dir, env, device):
    """加载训练好的策略"""
    print(f"\n加载模型: {model_dir}")

    # 加载配置
    config_path = os.path.join(os.path.dirname(os.path.dirname(model_dir)), 'params.json')
    if not os.path.exists(config_path):
        raise ValueError(f"找不到配置文件: {config_path}")

    with open(config_path, 'r') as f:
        saved_config = json.load(f)

    class Args:
        def __init__(self, config):
            for k, v in config.items():
                setattr(self, k, v)

    args = Args(saved_config)

    # 创建策略
    num_agents = env.num_players
    policies = []

    for agent_id in range(num_agents):
        policy = R_MAPPOPolicy(
            args,
            env.observation_space[0],
            env.observation_space[0],
            env.action_space[0],
            device=device
        )

        # 加载权重
        actor_path = os.path.join(model_dir, f'actor_agent{agent_id}.pt')
        actor_state_dict = torch.load(actor_path, map_location=device)
        policy.actor.load_state_dict(actor_state_dict)
        policy.actor.eval()

        policies.append(policy)
        print(f"✓ 加载 Agent {agent_id} 模型")

    return policies


def generate_episode_gif(env, policies, episode_id, output_path, args, device):
    """生成单个 episode 的 GIF"""
    print(f"\n生成 Episode {episode_id + 1} 的 GIF...")

    obs = env.reset()
    num_agents = env.num_players

    # 初始化 RNN states
    rnn_states = [torch.zeros((1, 1, 64)) for _ in range(num_agents)]
    masks = [torch.ones((1, 1)) for _ in range(num_agents)]

    frames = []
    total_reward = 0
    step = 0
    done = False

    while not done and step < args.max_steps:
        # 渲染当前帧
        try:
            frame = env.render(mode='rgb_array')
            if frame is not None:
                frames.append(frame)
        except Exception as e:
            print(f"Warning: 渲染失败 - {e}")
            # 如果渲染失败，尝试使用 human 模式
            env.render(mode='human')

        # 获取动作
        actions = []
        for agent_id in range(num_agents):
            obs_tensor = torch.FloatTensor(obs[agent_id]).unsqueeze(0).to(device)

            with torch.no_grad():
                action, rnn_state = policies[agent_id].act(
                    obs_tensor,
                    rnn_states[agent_id],
                    masks[agent_id],
                    available_actions=None,
                    deterministic=args.deterministic
                )

            actions.append(int(action.item()))
            rnn_states[agent_id] = rnn_state

        # 执行动作
        obs, rewards, dones, _ = env.step(actions)

        reward = rewards[0][0]
        total_reward += reward
        step += 1
        done = dones[0]

        if reward > 0:
            print(f"  步骤 {step}: 获得奖励 {reward}! (累计: {total_reward})")

        # 更新 masks
        masks = [torch.ones((1, 1)) if not dones[i] else torch.zeros((1, 1))
                 for i in range(num_agents)]

    print(f"  Episode 完成: {step} 步, 总奖励 {total_reward}")

    # 保存 GIF
    if frames and HAS_IMAGEIO:
        duration = 1.0 / args.fps  # 每帧持续时间（秒）
        imageio.mimsave(output_path, frames, duration=duration)
        print(f"  ✓ GIF 已保存: {output_path} ({len(frames)} 帧)")
        return True
    elif not HAS_IMAGEIO:
        print(f"  ✗ 无法保存 GIF: imageio 未安装")
        return False
    else:
        print(f"  ✗ 无法保存 GIF: 没有捕获到帧")
        return False


def main():
    args = parse_args()

    if not HAS_IMAGEIO:
        print("\n错误: 需要安装 imageio 库")
        print("请运行: pip install imageio")
        return

    # 设置设备
    device = torch.device("cuda:0" if args.cuda and torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)

    # 创建环境
    print(f"\n创建 Overcooked 环境: {args.layout_name}")
    env = Overcooked(
        layout_name=args.layout_name,
        max_timesteps=args.max_steps,
        obs_type='vector',
        multi_round=False
    )

    # 加载模型
    policies = load_policies(args.model_dir, env, device)

    # 生成 GIF
    print(f"\n{'='*60}")
    print(f"开始生成 {args.n_episodes} 个 episode 的 GIF 动画")
    print(f"输出目录: {args.output_dir}")
    print(f"帧率: {args.fps} FPS")
    print(f"{'='*60}")

    success_count = 0
    for episode_id in range(args.n_episodes):
        output_path = os.path.join(args.output_dir, f'episode_{episode_id + 1}.gif')

        try:
            if generate_episode_gif(env, policies, episode_id, output_path, args, device):
                success_count += 1
        except KeyboardInterrupt:
            print("\n\n用户中断")
            break
        except Exception as e:
            print(f"\n错误: 生成 Episode {episode_id + 1} 失败 - {e}")
            import traceback
            traceback.print_exc()

    print(f"\n{'='*60}")
    print(f"完成！成功生成 {success_count}/{args.n_episodes} 个 GIF")
    print(f"文件位置: {os.path.abspath(args.output_dir)}")
    print(f"{'='*60}\n")


if __name__ == "__main__":
    main()
